{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Approximate nearest neighbors in TSNE\n\n\nThis example presents how to chain KNeighborsTransformer and TSNE in a\npipeline. It also shows how to wrap the packages `annoy` and `nmslib` to\nreplace KNeighborsTransformer and perform approximate nearest neighbors.\nThese packages can be installed with `pip install annoy nmslib`.\n\nNote: Currently `TSNE(metric='precomputed')` does not modify the precomputed\ndistances, and thus assumes that precomputed euclidean distances are squared.\nIn future versions, a parameter in TSNE will control the optional squaring of\nprecomputed distances (see #12401).\n\nNote: In KNeighborsTransformer we use the definition which includes each\ntraining point as its own neighbor in the count of `n_neighbors`, and for\ncompatibility reasons, one extra neighbor is computed when\n`mode == 'distance'`. Please note that we do the same in the proposed wrappers.\n\nSample output::\n\n    Benchmarking on MNIST_2000:\n    ---------------------------\n    AnnoyTransformer:                    0.583 sec\n    NMSlibTransformer:                   0.321 sec\n    KNeighborsTransformer:               1.225 sec\n    TSNE with AnnoyTransformer:          4.903 sec\n    TSNE with NMSlibTransformer:         5.009 sec\n    TSNE with KNeighborsTransformer:     6.210 sec\n    TSNE with internal NearestNeighbors: 6.365 sec\n\n    Benchmarking on MNIST_10000:\n    ----------------------------\n    AnnoyTransformer:                    4.457 sec\n    NMSlibTransformer:                   2.080 sec\n    KNeighborsTransformer:               30.680 sec\n    TSNE with AnnoyTransformer:          30.225 sec\n    TSNE with NMSlibTransformer:         43.295 sec\n    TSNE with KNeighborsTransformer:     64.845 sec\n    TSNE with internal NearestNeighbors: 64.984 sec\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Author: Tom Dupre la Tour\n#\n# License: BSD 3 clause\nimport time\nimport sys\n\ntry:\n    import annoy\nexcept ImportError:\n    print(\"The package 'annoy' is required to run this example.\")\n    sys.exit()\n\ntry:\n    import nmslib\nexcept ImportError:\n    print(\"The package 'nmslib' is required to run this example.\")\n    sys.exit()\n\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib.ticker import NullFormatter\nfrom scipy.sparse import csr_matrix\n\nfrom sklearn.base import BaseEstimator, TransformerMixin\nfrom sklearn.neighbors import KNeighborsTransformer\nfrom sklearn.utils._testing import assert_array_almost_equal\nfrom sklearn.datasets import fetch_openml\nfrom sklearn.pipeline import make_pipeline\nfrom sklearn.manifold import TSNE\nfrom sklearn.utils import shuffle\n\nprint(__doc__)\n\n\nclass NMSlibTransformer(TransformerMixin, BaseEstimator):\n    \"\"\"Wrapper for using nmslib as sklearn's KNeighborsTransformer\"\"\"\n\n    def __init__(self, n_neighbors=5, metric='euclidean', method='sw-graph',\n                 n_jobs=1):\n        self.n_neighbors = n_neighbors\n        self.method = method\n        self.metric = metric\n        self.n_jobs = n_jobs\n\n    def fit(self, X):\n        self.n_samples_fit_ = X.shape[0]\n\n        # see more metric in the manual\n        # https://github.com/nmslib/nmslib/tree/master/manual\n        space = {\n            'sqeuclidean': 'l2',\n            'euclidean': 'l2',\n            'cosine': 'cosinesimil',\n            'l1': 'l1',\n            'l2': 'l2',\n        }[self.metric]\n\n        self.nmslib_ = nmslib.init(method=self.method, space=space)\n        self.nmslib_.addDataPointBatch(X)\n        self.nmslib_.createIndex()\n        return self\n\n    def transform(self, X):\n        n_samples_transform = X.shape[0]\n\n        # For compatibility reasons, as each sample is considered as its own\n        # neighbor, one extra neighbor will be computed.\n        n_neighbors = self.n_neighbors + 1\n\n        results = self.nmslib_.knnQueryBatch(X, k=n_neighbors,\n                                             num_threads=self.n_jobs)\n        indices, distances = zip(*results)\n        indices, distances = np.vstack(indices), np.vstack(distances)\n\n        if self.metric == 'sqeuclidean':\n            distances **= 2\n\n        indptr = np.arange(0, n_samples_transform * n_neighbors + 1,\n                           n_neighbors)\n        kneighbors_graph = csr_matrix((distances.ravel(), indices.ravel(),\n                                       indptr), shape=(n_samples_transform,\n                                                       self.n_samples_fit_))\n\n        return kneighbors_graph\n\n\nclass AnnoyTransformer(TransformerMixin, BaseEstimator):\n    \"\"\"Wrapper for using annoy.AnnoyIndex as sklearn's KNeighborsTransformer\"\"\"\n\n    def __init__(self, n_neighbors=5, metric='euclidean', n_trees=10,\n                 search_k=-1):\n        self.n_neighbors = n_neighbors\n        self.n_trees = n_trees\n        self.search_k = search_k\n        self.metric = metric\n\n    def fit(self, X):\n        self.n_samples_fit_ = X.shape[0]\n        metric = self.metric if self.metric != 'sqeuclidean' else 'euclidean'\n        self.annoy_ = annoy.AnnoyIndex(X.shape[1], metric=metric)\n        for i, x in enumerate(X):\n            self.annoy_.add_item(i, x.tolist())\n        self.annoy_.build(self.n_trees)\n        return self\n\n    def transform(self, X):\n        return self._transform(X)\n\n    def fit_transform(self, X, y=None):\n        return self.fit(X)._transform(X=None)\n\n    def _transform(self, X):\n        \"\"\"As `transform`, but handles X is None for faster `fit_transform`.\"\"\"\n\n        n_samples_transform = self.n_samples_fit_ if X is None else X.shape[0]\n\n        # For compatibility reasons, as each sample is considered as its own\n        # neighbor, one extra neighbor will be computed.\n        n_neighbors = self.n_neighbors + 1\n\n        indices = np.empty((n_samples_transform, n_neighbors),\n                           dtype=np.int)\n        distances = np.empty((n_samples_transform, n_neighbors))\n\n        if X is None:\n            for i in range(self.annoy_.get_n_items()):\n                ind, dist = self.annoy_.get_nns_by_item(\n                    i, n_neighbors, self.search_k, include_distances=True)\n\n                indices[i], distances[i] = ind, dist\n        else:\n            for i, x in enumerate(X):\n                indices[i], distances[i] = self.annoy_.get_nns_by_vector(\n                    x.tolist(), n_neighbors, self.search_k,\n                    include_distances=True)\n\n        if self.metric == 'sqeuclidean':\n            distances **= 2\n\n        indptr = np.arange(0, n_samples_transform * n_neighbors + 1,\n                           n_neighbors)\n        kneighbors_graph = csr_matrix((distances.ravel(), indices.ravel(),\n                                       indptr), shape=(n_samples_transform,\n                                                       self.n_samples_fit_))\n\n        return kneighbors_graph\n\n\ndef test_transformers():\n    \"\"\"Test that AnnoyTransformer and KNeighborsTransformer give same results\n    \"\"\"\n    X = np.random.RandomState(42).randn(10, 2)\n\n    knn = KNeighborsTransformer()\n    Xt0 = knn.fit_transform(X)\n\n    ann = AnnoyTransformer()\n    Xt1 = ann.fit_transform(X)\n\n    nms = NMSlibTransformer()\n    Xt2 = nms.fit_transform(X)\n\n    assert_array_almost_equal(Xt0.toarray(), Xt1.toarray(), decimal=5)\n    assert_array_almost_equal(Xt0.toarray(), Xt2.toarray(), decimal=5)\n\n\ndef load_mnist(n_samples):\n    \"\"\"Load MNIST, shuffle the data, and return only n_samples.\"\"\"\n    mnist = fetch_openml(data_id=41063)\n    X, y = shuffle(mnist.data, mnist.target, random_state=42)\n    return X[:n_samples], y[:n_samples]\n\n\ndef run_benchmark():\n    datasets = [\n        ('MNIST_2000', load_mnist(n_samples=2000)),\n        ('MNIST_10000', load_mnist(n_samples=10000)),\n    ]\n\n    n_iter = 500\n    perplexity = 30\n    # TSNE requires a certain number of neighbors which depends on the\n    # perplexity parameter.\n    # Add one since we include each sample as its own neighbor.\n    n_neighbors = int(3. * perplexity + 1) + 1\n\n    transformers = [\n        ('AnnoyTransformer', AnnoyTransformer(n_neighbors=n_neighbors,\n                                              metric='sqeuclidean')),\n        ('NMSlibTransformer', NMSlibTransformer(n_neighbors=n_neighbors,\n                                                metric='sqeuclidean')),\n        ('KNeighborsTransformer', KNeighborsTransformer(\n            n_neighbors=n_neighbors, mode='distance', metric='sqeuclidean')),\n        ('TSNE with AnnoyTransformer', make_pipeline(\n            AnnoyTransformer(n_neighbors=n_neighbors, metric='sqeuclidean'),\n            TSNE(metric='precomputed', perplexity=perplexity,\n                 method=\"barnes_hut\", random_state=42, n_iter=n_iter), )),\n        ('TSNE with NMSlibTransformer', make_pipeline(\n            NMSlibTransformer(n_neighbors=n_neighbors, metric='sqeuclidean'),\n            TSNE(metric='precomputed', perplexity=perplexity,\n                 method=\"barnes_hut\", random_state=42, n_iter=n_iter), )),\n        ('TSNE with KNeighborsTransformer', make_pipeline(\n            KNeighborsTransformer(n_neighbors=n_neighbors, mode='distance',\n                                  metric='sqeuclidean'),\n            TSNE(metric='precomputed', perplexity=perplexity,\n                 method=\"barnes_hut\", random_state=42, n_iter=n_iter), )),\n        ('TSNE with internal NearestNeighbors',\n         TSNE(metric='sqeuclidean', perplexity=perplexity, method=\"barnes_hut\",\n              random_state=42, n_iter=n_iter)),\n    ]\n\n    # init the plot\n    nrows = len(datasets)\n    ncols = np.sum([1 for name, model in transformers if 'TSNE' in name])\n    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, squeeze=False,\n                             figsize=(5 * ncols, 4 * nrows))\n    axes = axes.ravel()\n    i_ax = 0\n\n    for dataset_name, (X, y) in datasets:\n\n        msg = 'Benchmarking on %s:' % dataset_name\n        print('\\n%s\\n%s' % (msg, '-' * len(msg)))\n\n        for transformer_name, transformer in transformers:\n            start = time.time()\n            Xt = transformer.fit_transform(X)\n            duration = time.time() - start\n\n            # print the duration report\n            longest = np.max([len(name) for name, model in transformers])\n            whitespaces = ' ' * (longest - len(transformer_name))\n            print('%s: %s%.3f sec' % (transformer_name, whitespaces, duration))\n\n            # plot TSNE embedding which should be very similar across methods\n            if 'TSNE' in transformer_name:\n                axes[i_ax].set_title(transformer_name + '\\non ' + dataset_name)\n                axes[i_ax].scatter(Xt[:, 0], Xt[:, 1], c=y, alpha=0.2,\n                                   cmap=plt.cm.viridis)\n                axes[i_ax].xaxis.set_major_formatter(NullFormatter())\n                axes[i_ax].yaxis.set_major_formatter(NullFormatter())\n                axes[i_ax].axis('tight')\n                i_ax += 1\n\n    fig.tight_layout()\n    plt.show()\n\n\nif __name__ == '__main__':\n    test_transformers()\n    run_benchmark()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}